import tensorflow as tf
import cv2
import numpy as np

class ImportGraph():
    def __init__(self, model_name, loc):
        self.graph = tf.Graph()
        self.sess = tf.Session(graph=self.graph)
        self.lable = model_name

        with self.graph.as_default():
            # 从指定路径加载模型到局部图中
            saver = tf.train.import_meta_graph(loc + '.meta')
            saver.restore(self.sess, loc)

            self.in_x = self.sess.graph.get_tensor_by_name('input_x:0')
            self.y = self.sess.graph.get_tensor_by_name('output:0')

    def imgHandle(self, img):
        # 将原图像灰度化
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        # 平均滤波
        blur = cv2.blur(gray, (5, 5))
        # 简单阈值的二值化
        ret, thresh = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        # 查找轮廓
        contours, hierarchy = cv2.findContours(thresh, cv2.RETR_TREE, \
                                              cv2.CHAIN_APPROX_SIMPLE)[-2:]
        # 找最大面积轮廓
        area_max = 0
        _cnt = None
        for cnt in contours:
            area = cv2.contourArea(cnt)
            if area > area_max:
                area_max = area
                _cnt = cnt
        # 直边界矩形
        x, y, w, h = cv2.boundingRect(_cnt)
        # print(x, y, w, h)
        frame = img[y:h + y, x:w + x]
        return frame


    def run(self, frame):
        if self.lable == 'hanzi':
            name = "富", "强", "民", "主", "文", "明", "和", "谐"
            frame = self.imgHandle(frame.copy())
        elif self.lable == 'bzw':
            frame = frame.copy()[20:340, 160:480]
            name = "报警器", "道闸", "立体显示", "TFT", "无线充电", "语音播报"," "," "," "

        elif self.lable == 'jtbz':
            name = "停止让行", "禁止驶入", "向右转弯", "向左转弯", "禁止掉头", "禁止左拐", "车辆慢行", "右道封闭", "左道封闭"
            frame = self.imgHandle(frame.copy())

        img = frame.copy()
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = cv2.resize(img, (227, 227), interpolation=cv2.INTER_CUBIC)
        data = img.flatten()
        data = [np.asarray(data)]

        scores = self.sess.run(self.y, feed_dict={self.in_x: data})
        num = (np.argmax(scores, 1))
        print('index',np.ndarray.sum(num))
        lable = name[np.ndarray.sum(num)]
        img = putText(frame, lable, (0, 0), 'simsun.ttc', (0, 0, 255), 32)
        return lable, img


if __name__ == "__main__":
    hanzi_model = ImportGraph('hanzi', "./models/hanzi_model")
    bzw_model = ImportGraph('bzw', "./models/bzw_model")
    jtbz_model = ImportGraph("jtbz","./models/jtbz_model")

    img0 = cv2.imread('./0_1.jpg')
    img1 = cv2.imread('./test.jpg')
    img2 = cv2.imread('./test/06/215.jpg')

    cv2.imshow('img1', img1)
    cv2.imshow('img0', img0)
    cv2.waitKey(1)

    result, _ = hanzi_model.run(img0)
    cv2.imshow('hanzi', _)

    print(result)
    result, _ = bzw_model.run(img1)
    cv2.imshow('bzw', _)

    print(result)
    result, _ = jtbz_model.run(img2)
    cv2.imshow('jtbz', _)
    print(result)

    cv2.waitKey()
